rm(list = ls())
library(reshape2)
library(ggplot2)
library(magrittr)
library(stringr)
devtools::load_all()

# 设置c(x)
pf <- function(k, psi){
  cc <- psi[1] + psi[2]*log(k)+psi[3]*(log(k))^2
  return(cc)
}

# 误差损失函数
griderror <- function(k,psi){
  Utoday <- 1/pf(k,psi) # u'(c)
  nextK <- k^theta+(1-delta)*k-pf(k,psi) # k_t+1
  EnextU <- (theta*k^(theta-1)+1-delta)/pf(nextK,psi) # E[u'(c_{t+1})]
  FF <- sum((Utoday - beta*EnextU)^2) # F
  return(FF)
}

# 参数设置
beta <- .98
delta <-  .1
theta <- .36

# 状态变量格点化
k <- seq(1,6,0.06)

# policy function
psi <- optim(rep(0.1,3),griderror,k=k,method = 'BFGS')
opt_c <- pf(k,psi$par)
next_k <- k^theta + (1-delta)*k-opt_c

picdata <- data.frame(k=k, next_k = next_k)
ggplot(picdata,aes(x = k, y = next_k)) + geom_line() +
  geom_line(aes(y = k),linetype = 2) + theme_bw()
